// infections.mod

@# define num_sectors = 4

var 
    c l lambda lambdae lambdas lambdad lambdai Vi Vd Vs e e_by_c e_by_l
	in d s delta ;

@# for j in 1 : num_sectors
    var c@{j} l@{j} m@{j} C@{j} L@{j} M@{j} chi@{j}t mbar@{j} Vmbar@{j} P@{j} ;
@# endfor

varexo	
	ee_i0 ;

parameters 
    beta sigma N gamma deltabar phi kappa rho e0 eta  
    us ud psim ;

@# for j in 1 : num_sectors
    parameters ec@{j} el@{j} chibar@{j} Deltachi@{j} z@{j} ;
@# endfor

load input ;

beta			= params(1) ;
sigma			= params(2) ;
N				= params(3) ;
gamma			= params(4) ;
deltabar		= params(5) ;
phi				= params(6) ;
kappa			= params(7) ;
rho				= params(8) ;
e0				= params(9) ;
eta				= params(10) ;
us				= params(11) ;
ud				= params(12) ;
psim 			= params(13) ;

@# for j in 1 : num_sectors
ec@{j}			= params(14 + (@{j}-1)*5) ;
el@{j}			= params(15 + (@{j}-1)*5) ;
chibar@{j}		= params(16 + (@{j}-1)*5) ;
Deltachi@{j}	= params(17 + (@{j}-1)*5) ;
z@{j} 			= params(18 + (@{j}-1)*5) ;
@# endfor

model ;

    @# for j in 1 : num_sectors

    // cj FOC    
    z@{j} * (c@{j}/c)^(-1/sigma)*(1/c) = lambda*P@{j} + 2*lambdae*ec@{j}*C@{j} ;

    // lj FOC
    (l@{j})^eta = lambda*P@{j} - 2*lambdae*(1-m@{j})*el@{j}*(1-M@{j})*L@{j} ;

    // mj FOC
    lambda*P@{j}*chi@{j}t*m@{j} = 1/(1-d(-1)-kappa*in(-1)) * beta * Vmbar@{j}(+1) + 2*lambdae*el@{j}*l@{j}*(1-M@{j})*L@{j} ;
    
    // chijt and mbar
    chi@{j}t = chibar@{j}*(1-Deltachi@{j}*(1-exp(-mbar@{j}(-1)))) ;
    mbar@{j} = psim*mbar@{j}(-1) + m@{j} ;

    // Vmbar
    Vmbar@{j} = beta*psim*Vmbar@{j}(+1) + (1/2)*lambda*(1-d(-1)-kappa*in(-1))*P@{j}*chibar@{j}*Deltachi@{j}*exp(-mbar@{j}(-1))*m@{j}^2 ;

    // goods market
    c@{j} = (1-d(-1)-kappa*in(-1))*(l@{j}-chi@{j}t/2*m@{j}^2)/(1-d(-1)) ;

    // aggregates
    C@{j} = c@{j} ; //N*(1-d(-1)) * c@{j} ;
    L@{j} = l@{j} ; //N*(1-kappa*in(-1)-d(-1)) * l@{j} ;
    M@{j} = m@{j} ;
   
    @# endfor

    // exposure by consumption
    e_by_c =     
    @# for j in 1:num_sectors
      (1/@{num_sectors})*(1-d(-1))*ec@{j}*c@{j}*C@{j} + 
    @# endfor
    0 ;

    // exposure by labor
    e_by_l =     
    @# for j in 1:num_sectors
      (1/@{num_sectors})*(1-d(-1)-kappa*in(-1))*(1-m@{j})*el@{j}*l@{j}*L@{j}*(1-M@{j}) +
    @# endfor
    0 ;

    // consumption
    c^((sigma-1)/sigma) =     
    @# for j in 1:num_sectors
      (1/@{num_sectors})*z@{j}*c@{j}^((sigma-1)/sigma) + 
    @# endfor
    0 ;

    // labor
    l =     
    @# for j in 1:num_sectors
      (1/@{num_sectors})*(1/(1+eta))*l@{j}^(1+eta) + 
    @# endfor
    0 ;

    // normalize prices
    1 = 
    @# for j in 1:num_sectors
        (1/@{num_sectors})*P@{j}^(1-sigma) + 
    @# endfor
    0 ;

    // e FOC
    lambdae = (lambdai - lambdas)*gamma*in(-1)*s(-1) ;

    // i FOC
    lambdai = -beta*Vi(+1) ;
    
    // d FOC
    lambdad = -beta*Vd(+1) ;

    // s FOC
    lambdas = -beta*Vs(+1) ;

    // Vi
    Vi = kappa * l - us*kappa - 
    (1-rho)*lambdai - 
    gamma*e*s(-1)*(lambdai-lambdas) - 
    (delta*kappa + phi*kappa*exp(phi*kappa*in(-1))*kappa^2*in(-1))*(lambdad-lambdai) +
    @# for j in 1:num_sectors
        - (1/@{num_sectors})*kappa*lambda*P@{j}*(l@{j} - chi@{j}t/2 * m@{j}^2) + 
        (1/@{num_sectors})*kappa*lambdae*(1-m@{j})*el@{j}*l@{j}*L@{j}*(1-M@{j}) +
    @# endfor
    0 ;

    // Vd
    Vd = l - ud - log(c) - lambdad + 
    @# for j in 1:num_sectors
        - (1/@{num_sectors})*lambda*P@{j}*(l@{j} - chi@{j}t/2 * m@{j}^2 - c@{j}) + 
        (1/@{num_sectors})*lambdae*(ec@{j}*c@{j}*C@{j} + (1-m@{j})*el@{j}*l@{j}*L@{j}*(1-M@{j})) +
    @# endfor
    0 ;

    // Vs
    Vs = -lambdai*gamma*e*in(-1) - lambdas*(1-gamma*e*in(-1)) ;

    // e 
    e = e0 + e_by_c + e_by_l ;

    //*****************//
    // Contagion Block //
    //*****************//

    in = in(-1) + gamma * e * in(-1) * s(-1) - rho * in(-1) - delta * kappa * in(-1) + ee_i0  ;
    d = d(-1) + delta * kappa * in(-1)   ;
    s = s(-1) - gamma * e * in(-1) * s(-1) - ee_i0 ;    
    delta = 1*(deltabar + (exp(phi * kappa * in(-1))-1)) ;

end;


initval;

c				= xinitial(1) ;
l				= xinitial(2) ;

lambda			= xinitial(4) ;
lambdae			= xinitial(5) ;
lambdas			= xinitial(6) ;
lambdad			= xinitial(7) ;
lambdai			= xinitial(8) ;
Vi				= xinitial(9) ;
Vd				= xinitial(10) ;
Vs				= xinitial(11) ;
e				= xinitial(12) ;
e_by_c			= xinitial(13) ;
e_by_l			= xinitial(14) ;
in				= xinitial(15) ;
d				= xinitial(16) ;
s				= xinitial(17) ;
delta 			= xinitial(18) ;

@# for j in 1:num_sectors
c@{j}			= xinitial(19 + (@{j}-1)*10) ;
l@{j}			= xinitial(20 + (@{j}-1)*10) ;
m@{j}			= xinitial(21 + (@{j}-1)*10) ;
C@{j}			= xinitial(22 + (@{j}-1)*10) ;
L@{j}			= xinitial(23 + (@{j}-1)*10) ;
M@{j}			= xinitial(24 + (@{j}-1)*10) ;
chi@{j}t		= xinitial(25 + (@{j}-1)*10) ;
mbar@{j}		= xinitial(26 + (@{j}-1)*10) ;
Vmbar@{j}		= xinitial(27 + (@{j}-1)*10) ;
P@{j} 			= xinitial(28 + (@{j}-1)*10) ;
@# endfor

end;



resid 

endval ;

c				= xfinal(1) ;
l				= xfinal(2) ;

lambda			= xfinal(4) ;
lambdae			= xfinal(5) ;
lambdas			= xfinal(6) ;
lambdad			= xfinal(7) ;
lambdai			= xfinal(8) ;
Vi				= xfinal(9) ;
Vd				= xfinal(10) ;
Vs				= xfinal(11) ;
e				= xfinal(12) ;
e_by_c			= xfinal(13) ;
e_by_l			= xfinal(14) ;
in				= xfinal(15) ;
d				= xfinal(16) ;
s				= xfinal(17) ;
delta 			= xfinal(18) ;

@# for j in 1:num_sectors
c@{j}			= xfinal(19 + (@{j}-1)*10) ;
l@{j}			= xfinal(20 + (@{j}-1)*10) ;
m@{j}			= xfinal(21 + (@{j}-1)*10) ;
C@{j}			= xfinal(22 + (@{j}-1)*10) ;
L@{j}			= xfinal(23 + (@{j}-1)*10) ;
M@{j}			= xfinal(24 + (@{j}-1)*10) ;
chi@{j}t		= xfinal(25 + (@{j}-1)*10) ;
mbar@{j}		= xfinal(26 + (@{j}-1)*10) ;
Vmbar@{j}		= xfinal(27 + (@{j}-1)*10) ;
P@{j} 			= xfinal(28 + (@{j}-1)*10) ;
@# endfor

end ;

shocks;

    var ee_i0 ;
        periods 1 ;
        values (ee_i0_init) ;

end;


resid 

model_diagnostics;

simul(periods=300, maxit=10) ;
//forecast(periods=20);
